import glob
import os
import pathlib
import shutil
from multiprocessing.dummy.connection import Connection
from pathlib import Path
import cv2
import easydict

import random

classes = ["male,female"]
data_arrays = ["train", "valid", "test"]
data_weight = [7, 2, 1]


def create_folder_if_not_exit(folder):
    if os.path.exists(folder):
        return
    p = pathlib.Path(folder).parent
    if not os.path.exists(p):
        create_folder_if_not_exit(p)
    os.mkdir(folder)


class DataDivider:
    def __init__(self, datas=["train", "valid", "test"], weights=[7, 2, 1]):
        self.arr = []
        for i, d in enumerate(datas):
            self.arr.extend([d] * weights[i])
        random.shuffle(self.arr)
        print(self.arr)
        self.len = len(self.arr) - 1

    def get_random_group(self):
        # 按比例
        return self.arr[random.randint(0, self.len)]


# classes = [i for i in range(100)]
IMAGE_FOLDER = '/mnt/d/datasets/年龄性别数据集原数据/第一个/part2'
TARGET_FOLDER = f'/mnt/d/datasets/sex_and_age/age'


def get_sex(file_name: str):
    g = file_name.split("_")
    if len(g) > 2:
        return g[1]


def get_age(file_name: str):
    g = file_name.split("_")
    if len(g) > 1:
        return g[0]


def run(sex: str, p: Connection):
    images = glob.glob(f'{IMAGE_FOLDER}/*.jpg')
    dd = DataDivider()
    for im in images:
        if im.endswith('png'):
            continue
        img = cv2.imread(im)
        img_set = easydict.EasyDict({})
        img_set.img = img
        img_set.name = Path(im).stem
        img_set.sex = get_sex(img_set.name)
        img_set.age = get_age(img_set.name)
        p.send(img_set)
        sets = p.recv()
        # 保存
        group = dd.get_random_group()
        fname = Path(im).name
        img_folder = f'{TARGET_FOLDER}/{sex}/{group}/images'
        txt_folder = f'{TARGET_FOLDER}/{sex}/{group}/labels'
        create_folder_if_not_exit(img_folder)
        create_folder_if_not_exit(txt_folder)
        if sets is None or len(sets) > 1:
            continue
        with open(f'{txt_folder}/{img_set.name}.txt', 'w') as f:
            content: str = sets[0].rstrip()
            if content is None:
                continue
            if sex == 'age':
                content = f'{map_age_to_cls_index(img_set.age)} {content}'
            elif sex == 'sex':
                content = f'{img_set.sex} {content}'
            print(content)
            f.write(content)
        s = im
        t = f'{img_folder}/{fname}'
        shutil.copyfile(s, t)


# b = [i for i in range(5, 101, 10)]


def map_age_to_cls(age):
    ages = ['<5', '5<=x<15', '15<=x<25', '25<=x<35', '35<=x<45', '45<=x<55', '55<=x<65', '65<=x<75', '75<=x<85', '>=85']
    if age < 5:
        return ages[0]
    elif age >= 85:
        return ages[-1]
    else:
        return ages[(age - 5) // 10 + 1]


def map_age_to_cls_index(age):
    ages = ['<5', '5<=x<15', '15<=x<25', '25<=x<35', '35<=x<45', '45<=x<55', '55<=x<65', '65<=x<75', '75<=x<85',
            '>=85']
    age = int(age)
    if age < 5:
        return 0
    elif age >= 85:
        return len(ages) - 1
    else:
        return (age - 5) // 10 + 1


if __name__ == '__main__':
    head_detector = HeadDetector()

    run()

#  #文件名 ： 年龄_性别_种族_日期
#
# origin_dataset_folder = f'{FOLDER}/data_sets'
# train_folder = f'{FOLDER}/{result_folder}/train'
# train_images_folder = f'{FOLDER}/{result_folder}/train/images'
# train_labels_folder = f'{FOLDER}/{result_folder}/train/labels'
#
# valid_folder = f'{FOLDER}/{result_folder}/valid'
# valid_images_folder = f'{FOLDER}/{result_folder}/valid/images'
# valid_labels_folder = f'{FOLDER}/{result_folder}/valid/labels'
# result_folder = f'{FOLDER}/{result_folder}'
#
